from WindPy import w
from comm_tools.config import Config
from comm_tools.logger import log_progress
from comm_tools.database_mysql import open_mysql, load_to_MySQL_on_Cloud
from comm_tools import date_tool
import pandas as pd

code_list = ["000001.SH", "DJI.GI", "HSI.HI", "IXIC.GI", "000012.SH",
    "000159.SH", "H50069.CSI", "IC.CFE", "IH.CFE", "TF.CFE", "T.CFE",
    "N225.GI", "SPX.GI", "000300.SH", "IF.CFE"]

def format_to_pd(list_data, dates):
    df = pd.DataFrame(list_data)

    df.columns = code_list

    df['Date'] = dates

    return df

def add_trading_days(data_df, begin_date, end_date):
    days = w.tdayscount(begin_date, end_date).Data[0][0]
    data_df['Days'] = days

def add_aggr_mean_amt(data_df, begin_date, end_date):
    returned_data = w.wsd("000001.SH", "amt", begin_date, end_date, "") # 交易日

    # 检查数据是否成功获取
    if returned_data.ErrorCode != 0:
        log_progress("数据获取失败，请检查日期是否正确或网络连接。")

    df = pd.DataFrame(returned_data.Data, index=["Transaction Amount"], columns=returned_data.Times).T

    total_amount_yuan = df["Transaction Amount"].sum()
    mean_amount_yuan = df["Transaction Amount"].mean()

    total_amount_wan = total_amount_yuan / 1E8
    mean_amount_wan = mean_amount_yuan / 1E8

    data_df['AGG_AMT'] = total_amount_wan
    data_df['AGG_MEAN'] = mean_amount_wan


def get_data(begin_date, end_date, code):
    if not w.isconnected():
        w.start()  # 默认命令超时时间为120秒，如需设置超时时间可以加入waitTime参数，例如waitTime=60,即设置命令超时时间为60秒

    # w.isconnected()  # 判断WindPy是否已经登录成功
    # 获取指数
    returned_data = w.wsd(code, "close", end_date, end_date, "Days=Alldays") # 自然日

    # 检查数据是否成功获取
    if returned_data.ErrorCode == 0:
        log_progress(f"指数收盘价为：{returned_data.Data}")
        log_progress(f"对应的时间为：{returned_data.Times}")
    else:
        log_progress("数据获取失败，请检查日期是否正确或网络连接。")

    data_df = format_to_pd(returned_data.Data, returned_data.Times)
    add_trading_days(data_df, begin_date, end_date)
    add_aggr_mean_amt(data_df, begin_date, end_date)

    return data_df

def save_data(df_data, table_name):
    with open_mysql() as engine:
        load_to_MySQL_on_Cloud(df_data, engine, table_name)


def upload(end_date=date_tool.get_last_day_of_previous_month()):
    # e.g. '2025-01-31'
    begin_date = date_tool.get_begin_date(end_date)
    data = get_data(begin_date, end_date, code_list)
    log_progress(data.to_string())

    c = Config()
    save_data(data, c.table_index_b)